import numpy as np
from urdfenvs.urdf_common.urdf_env import UrdfEnv
from bicycle_model_self import BicycleModelSelf
from urdfenvs.scene_examples.obstacles import *
from mpscenes.obstacles import *
from Create_Environment import Generate_Walls, Generate_Obstacles, Generate_Obstacle_Centers, create_grid_map, Generate_GoalPos, update_vehicle, initialize_plot
import matplotlib.pyplot as plt
from Simple_FreeSpace import compute_free_space
from RRT_Star import *
from A_Star_Planning import plan_path_through_goals
import pandas as pd

def run_prius(pos_idx, n_steps=200000, render=False, goal=True, obstacles=True, visualize=True):
    robots = [
        BicycleModelSelf(
            urdf='prius.urdf',
            mode="vel",
            scaling=1,
            wheel_radius=0.31265,
            wheel_distance=2.86,
            spawn_offset=np.array([0.0, 0.0, 0.0]),
            actuated_wheels=['front_right_wheel_joint', 'front_left_wheel_joint', 'rear_right_wheel_joint', 'rear_left_wheel_joint'],
            steering_links=['front_right_steer_joint', 'front_left_steer_joint'],
            facing_direction='-x'
        )
    ]

    env: UrdfEnv = UrdfEnv(dt=0.01, robots=robots, render=render)
    if pos_idx == 0:
        file_name = "2025-01-09_13_30_43.pkl"
    elif pos_idx == 4:
        file_name = "2025-01-09_13_33_04.pkl"
    else:
        file_name = "2025-01-09_13_35_54.pkl"

    # Define boundary and margin for obstacles
    length = 50
    width = 80
    margin = 2.5
    Walls, Wall_Dict = Generate_Walls(length, width, Type='Parking Lot')
    for Wall in Walls:
        env.add_obstacle(Wall)

    # Define obstacles
    Obs_centers = Generate_Obstacle_Centers(length, width, n_obstacles=16, Type='Parking Lot')
    Obstacles , Obs_Dict = Generate_Obstacles(Obs_centers, Type='Parking Lot')
    for Obstacle in Obstacles:
        env.add_obstacle(Obstacle)

    # Integrate boundary and obstacles
    obstacles_list = []
    for obs in Obs_Dict:
        obstacles_list.append(obs)
    for walls in Wall_Dict:
        obstacles_list.append(walls)
    global_map = {'l': length, 'w': width, 'm': margin, 'o': obstacles_list}

    # Define Grid Resolution
    grid_size = 0.1
    grid_map, min_x, max_x, min_y, max_y = create_grid_map(Wall_Dict, Obs_Dict, grid_size=grid_size)

    freespace_map = compute_free_space(
        grid_map, grid_size, 2.34+1.42, 2.42-1.42, 0.7, min_x, max_x, min_y, max_y, num_theta=16)
    if visualize:   
        fig, ax = initialize_plot(freespace_map, min_x, max_x, min_y, max_y)
    # Define first action
    action = np.array([0.0, 0.0])  # Adjust action as necessary
    
    # Define start and goal in the world frame
    goals, goal_index = Generate_GoalPos(pos_idx)
    # pos0 = np.array([-20, -35, np.pi / 4, 0.0])
    if goal_index in {0, 2, 5, 6, 8}:
        pos0 = np.array([-20, -35, np.pi / 2, 0.0])
    else:
        pos0 = np.array([-20, -35, 0.0, 0.0])   
    goal = np.array([goals[-1, 0], goals[-1, 1], goals[-1, 2], 0.0])
    print(f'Pos {pos0}, index {goal_index}')
    
    if visualize:  
        update_vehicle(ax, goal[0:3])
        update_vehicle(ax, pos0[0:3])
    
    # Reset simulation environment
    ob = env.reset(pos=pos0[0:3])
    print(ob)
    # ---------Planning---------
    planning_success = False
    path_astar = plan_path_through_goals(pos0, goals, freespace_map, min_x, max_x, min_y, max_y, grid_size, ax, visualize=visualize)
    
    df = pd.read_pickle(file_name)
    origin_nodes_list = df['origin_nodes_list'].tolist()
    refined_nodes_list = df['refined_nodes_list'].tolist()
    explored_x = df['explored_x'].tolist()
    explored_y = df['explored_y'].tolist()
    path_x_original = []
    path_y_original = []
    path_x_refined = []
    path_y_refined = []
    for i in range(len(origin_nodes_list) - 1):
        current_node = origin_nodes_list[i]
        next_node = origin_nodes_list[i + 1]
        print(next_node)
        if isinstance(current_node, float) and np.isnan(current_node):
            break
        if isinstance(next_node, float) and np.isnan(next_node):
            break
        sx, sy, syaw, sk = current_node[0], current_node[1], current_node[2], current_node[3]
        gx, gy, gyaw, gk = next_node[0], next_node[1], next_node[2], next_node[3]

        path_x, path_y, _, path_length, _ = clothoid_path(sx, sy, syaw, sk, gx, gy, gyaw, gk)
        path_x_original.append(path_x)
        path_y_original.append(path_y)

    for i in range(len(refined_nodes_list) - 1):
        current_node = refined_nodes_list[i]
        next_node = refined_nodes_list[i + 1]
        if isinstance(current_node, float) and np.isnan(current_node):
            break
        if isinstance(next_node, float) and np.isnan(next_node):
            break
        sx, sy, syaw, sk = current_node[0], current_node[1], current_node[2], current_node[3]
        gx, gy, gyaw, gk = next_node[0], next_node[1], next_node[2], next_node[3]
        path_x, path_y, _, path_length, _ = clothoid_path(sx, sy, syaw, sk, gx, gy, gyaw, gk)
        path_x_refined.append(path_x)
        path_y_refined.append(path_y)
    
    plot_pkl(ax, explored_x, explored_y, path_x_refined, path_y_refined, path_astar)
    return planning_success


def plot_pkl(ax, explored_x, explored_y, refined_x, refined_y, Astar):
    """
    Compare and visualize the explored points and refined path on an existing plot,
    and save the plot to a specified folder.

    Parameters:
    ax -- Matplotlib axis object for plotting.
    explored_x, explored_y -- Coordinates of the explored points (lists of x and y points).
    refined_x, refined_y -- Coordinates of the refined path (lists of x and y points).
    Astar -- Boolean indicating whether the visualization is for A* algorithm.

    Returns:
    None. The plot is updated on the given axis and saved as a PNG file.
    """
    # Plot the explored points with dotted line and 'o' markers
    if explored_x and explored_y:
        ax.scatter(explored_x, explored_y, color='#48597e', s=1, marker='o', label="Explored Points")
    else:
        print("No Explored Points data provided!")

    # Plot the refined path with solid line and 'x' markers
    if refined_x and refined_y:
        ax.scatter(refined_x, refined_y, color='red', s=0.5, marker='x', label="Refined Path Points")
    else:
        print("No Refined Path data provided!")

    # Add information about A* algorithm if applicable
    if Astar:
        ax.set_title("A* Algorithm Visualization", fontsize=16)
    else:
        ax.set_title("Path Visualization", fontsize=16)

    # Add a legend
    ax.legend(fontsize=12)
    # Show the plot
    plt.show(block=True)

if __name__ == "__main__":
    pos_indices = [0, 4, 6]
    for i in pos_indices:
        for _ in range(1):
            _ = run_prius(pos_idx=i, render=False, visualize=True)


